#!/usr/bin/env python3
import os
import sys
import time
from collections import Counter, deque
from typing import Deque, Optional, Tuple

import rospy
from std_msgs.msg import String


def _resolve_paths() -> Tuple[str, str, str]:
    # scripts -> dog_motion -> src -> nav_ws -> repo_root
    repo_root = os.path.dirname(
        os.path.dirname(
            os.path.dirname(
                os.path.dirname(
                    os.path.dirname(os.path.abspath(__file__))
                )
            )
        )
    )
    detect_dir = os.path.join(repo_root, 'yolo-jetson', 'detect')
    rs2_dir = os.path.join(repo_root, 'yolo-jetson', 'rs2')
    return repo_root, detect_dir, rs2_dir


REPO_ROOT, DETECT_DIR, RS2_DIR = _resolve_paths()
if DETECT_DIR not in sys.path:
    sys.path.append(DETECT_DIR)
if RS2_DIR not in sys.path:
    sys.path.append(RS2_DIR)

from yolo_pipe_ocr_hsv_optimal import create_ocr_hsv_pipeline  # noqa: E402
from GStreamerWrapper import GStreamerWrapper  # noqa: E402


def majority_vote(cache: Deque[Tuple[str, str, str]]) -> Optional[Tuple[str, str, str]]:
    """使用多数投票法从缓存中确定最终结果"""
    if len(cache) == 0:
        return None
    merged = ['%s-%s-%s' % t for t in cache]
    top, _ = Counter(merged).most_common(1)[0]
    c, n, d = top.split('-')
    return c, n, d


class YOLOOcrHsvFinalNode:
    """改进后的YOLO OCR HSV节点，基于yolo_pipe_ocr_hsv_final.py"""
    
    def __init__(self) -> None:
        rospy.init_node('yolo_ocr_hsv_final_node', anonymous=False)

        # 基础参数
        self.width = rospy.get_param('~width', 1920)
        self.height = rospy.get_param('~height', 1080)
        self.first_stage_conf = rospy.get_param('~first_stage_conf', 0.6)
        self.final_conf = rospy.get_param('~final_conf', 0.6)
        self.top_cut_ratio = rospy.get_param('~top_cut_ratio', 0.25)
        self.width_center_ratio = rospy.get_param('~width_center_ratio', 0.5)
        self.cache_size = int(rospy.get_param('~cache_size', 10))
        self.target_fps = int(rospy.get_param('~target_fps', 12))
        self.debug_mode = rospy.get_param('~debug', False)

        # 模型路径发现
        candidate_models = [
            os.path.join(DETECT_DIR, 'weights', 'all_yolo_all_epochs_0811', 'best.pt'),
            os.path.join(DETECT_DIR, 'weights', 'all_yolo_all_epochs_0808', 'best.pt'),
            os.path.join(DETECT_DIR, 'weights', 'all_yolo_all_epochs_0805', 'best.pt'),
            os.path.join(DETECT_DIR, 'weights', 'all_yolo_all_epochs_0804', 'best.pt'),
        ]
        self.model_path = rospy.get_param('~model_path', '')
        if not self.model_path:
            for p in candidate_models:
                if os.path.exists(p):
                    self.model_path = p
                    rospy.loginfo(f'Found model: {p}')
                    break
        if not self.model_path or not os.path.exists(self.model_path):
            rospy.logerr('No model weights found. Set ~model_path param to a valid .pt file.')
            raise RuntimeError('model_path not found')

        # 初始化组件
        self.pipeline = create_ocr_hsv_pipeline(
            model_path=self.model_path, 
            debug=self.debug_mode
        )
        self.gs = GStreamerWrapper(width=self.width, height=self.height)

        # 启动时预热引擎
        self._warmup_engine()

        # ROS 发布/订阅
        self.result_pub = rospy.Publisher('/yolo_result', String, queue_size=10)
        self.cmd_sub = rospy.Subscriber('/yolo_detect', String, self._cmd_cb, queue_size=10)
        
        # 新增：发布详细结果
        self.detailed_result_pub = rospy.Publisher('/yolo_detailed_result', String, queue_size=10)

        # 状态管理
        self._running = False
        self._cache: Deque[Tuple[str, str, str]] = deque(maxlen=self.cache_size)
        self._frame_interval = 1.0 / max(1, self.target_fps)
        self._last_time = 0.0
        self._frame_count = 0
        self._success_count = 0

        # 定时器循环处理
        self._timer = rospy.Timer(rospy.Duration(0.01), self._on_timer)

        rospy.loginfo('yolo_ocr_hsv_final_node ready. Send "yolostart" on /yolo_detect to begin.')
        rospy.loginfo(f'Model: {self.model_path}')
        rospy.loginfo(f'Resolution: {self.width}x{self.height}')
        rospy.loginfo(f'Target FPS: {self.target_fps}')

    def _warmup_engine(self) -> None:
        """预热YOLO引擎以确保TensorRT引擎完全加载"""
        start_time = time.time()
        frame = None
        # 尝试获取真实帧进行预热
        while time.time() - start_time < 3.0:
            frame = self.gs.GetFrame()
            if frame is not None:
                break
            time.sleep(0.05)
        if frame is None:
            try:
                import numpy as np
                frame = np.zeros((self.height, self.width, 3), dtype=np.uint8)
                rospy.logwarn('Using dummy frame for warm-up')
            except Exception as e:
                rospy.logerr(f'Failed to allocate dummy frame for warm-up: {e}')
                raise
        try:
            _ = self.pipeline.process_image(
                frame,
                top_cut_ratio=self.top_cut_ratio,
                width_center_ratio=self.width_center_ratio,
                first_stage_conf=self.first_stage_conf,
                final_conf=self.final_conf,
            )
            rospy.loginfo('YOLO engine warm-up completed.')
        except Exception as e:
            rospy.logerr(f'YOLO engine warm-up failed: {e}')
            raise

    def _cmd_cb(self, msg: String) -> None:
        """处理命令消息"""
        cmd = (msg.data or '').strip().lower()
        if cmd == 'yolostart' and not self._running:
            self._cache.clear()
            self._running = True
            self._frame_count = 0
            self._success_count = 0
            rospy.loginfo('YOLO OCR+HSV Final started')
        elif cmd in ('stop', 'yolostop') and self._running:
            self._running = False
            self._cache.clear()
            rospy.loginfo('YOLO OCR+HSV Final stopped')
        elif cmd == 'status':
            rospy.loginfo(f'Status: running={self._running}, frames={self._frame_count}, success={self._success_count}')

    def _on_timer(self, _evt) -> None:
        """定时器回调函数，处理图像帧"""
        if not self._running or rospy.is_shutdown():
            return
            
        frame = self.gs.GetFrame()
        if frame is None:
            return
            
        now = time.time()
        if now - self._last_time < self._frame_interval:
            return
        self._last_time = now
        self._frame_count += 1

        try:
            res = self.pipeline.process_image(
                frame,
                top_cut_ratio=self.top_cut_ratio,
                width_center_ratio=self.width_center_ratio,
                first_stage_conf=self.first_stage_conf,
                final_conf=self.final_conf,
            )

            cone = res.get('cone_color', 'None')
            num = res.get('number', 'None')
            dash = res.get('dashboard_state', 'None')

            # 记录有效结果
            if cone != 'None' and num != 'None' and dash != 'None':
                self._cache.append((cone, num, dash))
                self._success_count += 1
                rospy.loginfo(f'Valid: {cone}-{num}-{dash} ({len(self._cache)}/{self.cache_size})')

            # 发布详细结果（包含处理信息）
            detailed_msg = {
                'frame_count': self._frame_count,
                'success_count': self._success_count,
                'cache_size': len(self._cache),
                'cone_color': cone,
                'number': num,
                'dashboard_state': dash,
                'processing_info': res.get('processing_info', {}),
                'timestamp': now
            }
            self.detailed_result_pub.publish(String(data=str(detailed_msg)))

            # 当缓存满时进行多数投票
            if len(self._cache) >= self.cache_size:
                final_triplet = majority_vote(self._cache)
                if final_triplet is not None:
                    fc, fn, fd = final_triplet
                    msg = f'{fc}-{fn}-{fd}'
                    self.result_pub.publish(String(data=msg))
                    rospy.loginfo(f'Published /yolo_result: {msg}')
                    rospy.loginfo(f'Processing complete: {self._frame_count} frames, {self._success_count} successful')
                self._cache.clear()
                self._running = False

        except Exception as e:
            rospy.logerr(f'Error processing frame {self._frame_count}: {e}')


if __name__ == '__main__':
    try:
        node = YOLOOcrHsvFinalNode()
        rospy.spin()
    except Exception as e:
        rospy.logerr(f'yolo_ocr_hsv_final_node failed: {e}')
        import traceback
        traceback.print_exc()
